package paddledjl;

import ai.djl.inference.Predictor;
import ai.djl.modality.Classifications;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.ImageFactory;
import ai.djl.modality.cv.output.DetectedObjects;
import cn.hutool.core.date.ChineseDate;
import cn.hutool.core.date.DatePattern;
import cn.hutool.core.date.DateUtil;
import cn.hutool.core.date.Week;
import cn.hutool.core.io.FileUtil;
import lombok.extern.slf4j.Slf4j;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.InputStream;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.Date;
import java.util.List;
import java.util.StringJoiner;

/**
 * @author 小帅丶
 * @className PaddleOCRDJLSample
 * @Description PaddleOCR DJL示例
 * @Date 2021-12-13-15:00
 **/
@Slf4j
public class PaddleOCRDJLSample {

    private static String IMG_PATH = "F:\\testimg\\idcard\\idcard2.jpg";

    private static String NEW_IMG_PATH = "F:\\testmodel\\img\\";



    public static void main(String[] args) throws Exception {
        long startTime = System.currentTimeMillis();
        StringJoiner words = new StringJoiner("\n");
        //模型版本
        DJLPaddleOCRConfig.MODEL_VERSION version = DJLPaddleOCRConfig.MODEL_VERSION.DEFAULT;
        //加载检测、方向、识别模型
        Predictor<Image, Classifications> direction = DJLPaddleOCRConfig.getDirectionModel(version);
        Predictor<Image, String> recognition = DJLPaddleOCRConfig.getRecognitionModel(version);
        Predictor<Image, DetectedObjects> detector = DJLPaddleOCRConfig.getDetectionModel(version);

        //加载图片
        Path imageFile = Paths.get(IMG_PATH);
        Image img = ImageFactory.getInstance().fromFile(imageFile);

        //获取选择的置信度
        System.out.println(direction.predict(img));

        //获取认为是文字的区域进行框选
        DetectedObjects detectedObj = detector.predict(img);
        List<DetectedObjects.DetectedObject> boxes = detectedObj.items();

        //画边框
        Image newImage = img.duplicate();
        newImage.drawBoundingBoxes(detectedObj);
        BufferedImage wrappedImage = (BufferedImage) newImage.getWrappedImage();

        ByteArrayOutputStream os = new ByteArrayOutputStream();
        ImageIO.write(wrappedImage, "png", os);
        InputStream input = new ByteArrayInputStream(os.toByteArray());
        FileUtil.writeFromStream(input, NEW_IMG_PATH + System.currentTimeMillis() + ".png");

        for (DetectedObjects.DetectedObject box : boxes) {
            Image sample = PaddleOCRUtil.getSubImage(img, box.getBoundingBox());
            //识别的文字
            words.add(recognition.predict(sample));
        }
        long endTime = System.currentTimeMillis();
        System.out.println("识别的内容:"+words.toString());
        System.out.println("总耗时:"+(endTime-startTime));
    }

}
